# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
下标访问发射器

处理 IRSubscript 节点，生成数组/容器的下标访问代码。
"""

from ...ir import IRNode
from ...ir.expressions.access import IRSubscript, IRSlice
from ..core.emitter import NodeTypeEmitter, EmitResult
from ..core.context import CodeGenContext
from ..configs.operators import get_operator_precedence
from ..core.generator import get_global_generator
from mcpy.symbol_meta import TypeKind


class SubscriptEmitter(NodeTypeEmitter):
    """下标访问发射器

    处理各种下标访问：
    - 数组访问：arr[0] -> arr[0]
    - 容器访问：vec[i] -> vec[i]
    - 字典访问：dict[key] -> dict[key]
    """

    def __init__(self):
        super().__init__([IRSubscript])

    def emit(self, node: IRSubscript, context: CodeGenContext) -> EmitResult:
        """生成下标访问代码

        策略：利用类型配置中的 __getitem__ 方法配置
        - 如果是切片，使用切片专用函数
        - 如果类型配置了 __getitem__，使用配置的模板生成代码
        - 否则使用标准的 [] 下标访问
        """
        import ast

        generator = get_global_generator()
        subscript_precedence = get_operator_precedence("subscript")

        # 生成被访问的对象表达式
        obj_code = generator.generate_expression(
            node.value, context, parent_precedence=subscript_precedence
        )

        # 检查是否是切片表达式
        if self._is_slice(node.slice):
            # 处理切片：list[start:stop:step]
            slice_code = self._emit_slice(node, obj_code, context)
            if slice_code:
                return EmitResult(code=slice_code, precedence=subscript_precedence)

        # 尝试使用类型配置中的 __getitem__ 方法
        getitem_code = self._try_emit_getitem_method(node, obj_code, context)
        if getitem_code:
            return EmitResult(code=getitem_code, precedence=subscript_precedence)

        # 如果没有找到 __getitem__ 配置，生成普通下标表达式
        index_code = generator.generate_expression(
            node.slice, context, parent_precedence=100
        )

        # 没有找到配置，使用标准的 [] 下标访问
        code = f"{obj_code}[{index_code}]"
        return EmitResult(code=code, precedence=subscript_precedence)

    def _try_emit_getitem_method(
        self, node: IRSubscript, obj_code: str, context: CodeGenContext
    ) -> str | None:
        """尝试使用类型符号的 __getitem__ 方法生成代码

        通过符号体系查找 __getitem__ 方法，支持 builtin 类型、用户自定义类和 stub

        Args:
            node: 下标访问节点
            obj_code: 对象表达式代码
            context: 代码生成上下文

        Returns:
            str | None: 生成的代码，如果无法使用配置则返回 None
        """
        from ..core.generator import get_global_generator

        generator = get_global_generator()
        # 从推断类型获取 TypeInfo
        type_info = node.value.get_extension_param("inferred_type")
        if not type_info:
            return None

        # 获取索引的类型信息用于方法重载匹配
        index_type = node.slice.get_extension_param("inferred_type")
        args_types = [index_type] if index_type else []

        # 查找 __getitem__ 方法符号
        method_symbol = None
        if context.current_node:
            method_symbol = context.current_node.lookup_symbol(
                "__getitem__",
                owner_type=type_info,
                func_args=args_types,
            )

        if not method_symbol:
            return None

        # 获取匹配的签名（考虑所有重载）
        # 对于 tuple，可能有常量索引和非常量索引两个版本
        matching_signature = method_symbol.find_matching_signature(
            [type_info] + args_types
        )

        # 如果找到了匹配的签名，还需要检查是否适合当前场景
        # 特别处理：弱类型 tuple (mc::variants) 不支持常量索引模板版本
        from ...ir.expressions.basic import IRConstant

        is_constant_index = isinstance(node.slice, IRConstant) and isinstance(
            node.slice.value, int
        )

        # 获取对象的 C++ 类型，判断是否是弱类型 tuple
        obj_cpp_type = (
            context.get_cpp_type(node.value) if hasattr(node, "value") else None
        )
        is_weak_tuple = obj_cpp_type == "mc::variants"

        # 如果匹配的签名要求常量索引，但对象是弱类型 tuple，需要切换
        if matching_signature:
            sig_constant_params = getattr(matching_signature, "constant_params", [])
            requires_constant = "index" in sig_constant_params

            if requires_constant and is_weak_tuple:
                # 弱类型 tuple 不支持常量索引版本，查找非常量索引版本
                matching_signature = None
                for sig in method_symbol.signatures:
                    sig_constant_params = getattr(sig, "constant_params", [])
                    if "index" not in sig_constant_params:
                        matching_signature = sig
                        break

        if not matching_signature:
            # 如果没有找到匹配的签名，尝试选择最合适的签名
            # 优先选择：非常量索引版本（更通用）
            for sig in method_symbol.signatures:
                sig_constant_params = getattr(sig, "constant_params", [])
                if "index" not in sig_constant_params:
                    matching_signature = sig
                    break

            # 如果还是没有找到，使用第一个签名
            if not matching_signature and method_symbol.signatures:
                matching_signature = method_symbol.signatures[0]
            else:
                return None

        # 获取模板
        template = matching_signature.emit_template
        if not template:
            return None

        # 替换模板中的类型参数（统一使用 template_emitter 的逻辑）
        template = self._replace_template_parameters(
            template, type_info, method_symbol, matching_signature, node, context
        )

        if not template:
            return None

        # 添加必要的包含文件
        includes = getattr(matching_signature, "includes", [])
        for include in includes:
            context.require_include(include)

        # 获取索引参数的期望类型
        expected_index_cpp_type = self._get_index_expected_type(
            method_symbol, matching_signature, node, context
        )

        # 生成索引表达式
        index_context = context.create_child_context(
            expected_cpp_type=expected_index_cpp_type
        )
        index_code = generator.generate_expression(
            node.slice, index_context, parent_precedence=100
        )

        # 应用类型转换（如果需要：mc::variant -> int）
        if expected_index_cpp_type:
            actual_index_cpp_type = context.get_cpp_type(node.slice)
            index_code = self._apply_type_conversion(
                index_code, actual_index_cpp_type, expected_index_cpp_type, context
            )

        # 检查是否需要常量索引（通过 constant_params 列表）
        requires_constant_index = False
        if (
            hasattr(matching_signature, "constant_params")
            and matching_signature.constant_params
        ):
            # 检查 index 参数是否在常量参数列表中
            if "index" in matching_signature.constant_params:
                requires_constant_index = True

        if requires_constant_index:
            # 需要常量索引版本
            if is_constant_index:
                # 对于常量索引，直接使用数值（不带引号）
                # 已经通过 is_constant_index 检查，确保 node.slice 是 IRConstant
                constant_slice = (
                    node.slice if isinstance(node.slice, IRConstant) else None
                )
                if constant_slice and isinstance(constant_slice.value, int):
                    constant_index = str(constant_slice.value)
                    code = template.replace("{self}", obj_code).replace(
                        "{index}", constant_index
                    )
                else:
                    # 无法提取常量值，回退
                    return None
            else:
                # 需要常量索引但提供的是变量，尝试查找非常量索引版本
                for sig in method_symbol.signatures:
                    if sig == matching_signature:
                        continue  # 跳过当前匹配的签名
                    sig_constant_params = getattr(sig, "constant_params", [])
                    if "index" not in sig_constant_params:
                        # 找到非常量索引版本，使用它
                        matching_signature = sig
                        template = sig.emit_template
                        requires_constant_index = False
                        # 重要：切换模板后需要重新进行模板参数替换
                        template = self._replace_template_parameters(
                            template,
                            type_info,
                            method_symbol,
                            matching_signature,
                            node,
                            context,
                        )
                        if not template:
                            return None
                        break

                # 如果还是需要常量索引但没有常量值，回退到标准访问
                if requires_constant_index:
                    return None
                else:
                    # 切换到非常量索引版本，生成代码
                    code = template.replace("{self}", obj_code).replace(
                        "{index}", index_code
                    )
        else:
            # 替换模板占位符（非常量索引版本）
            code = template.replace("{self}", obj_code).replace("{index}", index_code)

        return code

    def _get_index_expected_type(
        self, method_symbol, signature, node: IRSubscript, context: CodeGenContext
    ) -> str | None:
        """从方法符号的 TypeInfo 中获取索引参数的期望 C++ 类型

        Args:
            method_symbol: 方法符号对象
            signature: 方法签名对象
            node: 下标访问节点
            context: 代码生成上下文

        Returns:
            str | None: 期望的 C++ 类型，如果无法确定则返回 None
        """
        # 优先从 method_symbol 的 TypeInfo 获取（已经进行了模板参数替换）
        if method_symbol and method_symbol.type_info and method_symbol.type_info.args:
            from mcpy.type_system import get_cpp_type

            # 方法的 TypeInfo.args 包含所有参数类型：[self, index, ...]
            # 索引参数是第二个（index=1）
            if len(method_symbol.type_info.args) >= 2:
                index_param_type_info = method_symbol.type_info.args[1]
                try:
                    cpp_type = get_cpp_type(index_param_type_info, node)
                    return cpp_type
                except:
                    pass

        return None

    def _replace_template_parameters(
        self,
        template: str,
        type_info,
        method_symbol,
        matching_signature,
        node: IRSubscript,
        context: CodeGenContext,
    ) -> str:
        """替换模板中的类型参数占位符

        统一处理模板参数替换逻辑，包括：
        1. 替换 owner_type 的类型参数（如 tuple[T, Args...] 中的 {T}, {Args...}）
        2. 替换返回类型中的模板参数占位符（如 -> T 中的 {T}）

        Args:
            template: 原始模板字符串
            type_info: owner_type 的 TypeInfo
            method_symbol: 方法符号对象
            matching_signature: 匹配的签名对象
            node: 下标访问节点
            context: 代码生成上下文

        Returns:
            替换后的模板字符串
        """
        from mcpy.codegen.emitters.template_emitter import TemplateEmitter
        from mcpy.type_system import get_cpp_type

        # 处理返回类型占位符（需要在模板参数替换之前，避免被错误替换）
        # 例如：tuple_getitem<{T}>(...) 中的 {T} 是返回类型占位符，不是模板参数
        # 策略：先用临时占位符替换返回类型的 {T}，避免被步骤1替换，然后再替换回实际类型
        return_type_temp_placeholder = "__PLACEHOLDER__"
        return_type_cpp_type = None

        if matching_signature.return_type:
            return_type_str = matching_signature.return_type.strip()
            # 检查返回类型是否是简单的模板参数名（如 "T"），且在模板中作为返回类型占位符存在
            # 支持两种格式：
            # 1. 单个模板参数：<{T}>
            # 2. 多个模板参数：<{index}, {T}> 或 <{T}, ...>
            return_type_placeholder_pattern = f"{{{return_type_str}}}"
            if return_type_str and return_type_placeholder_pattern in template:
                # 优先从 node 的 inferred_type 获取实际类型（支持负数索引等特殊情况）
                inferred_type = node.get_extension_param("inferred_type")
                if inferred_type:
                    return_type_cpp_type = get_cpp_type(inferred_type, node)
                else:
                    # 回退：从 method_symbol 获取（已通过符号系统替换了模板参数）
                    if (
                        method_symbol
                        and method_symbol.type_info
                        and method_symbol.type_info.return_type
                    ):
                        return_type_info = method_symbol.type_info.return_type
                        return_type_cpp_type = get_cpp_type(return_type_info, node)

                if return_type_cpp_type:
                    # 在步骤1之前，用临时占位符替换返回类型的 {T}，避免被步骤1替换
                    # 支持多个模板参数的情况，只替换 {T} 部分
                    template = template.replace(
                        return_type_placeholder_pattern, return_type_temp_placeholder
                    )

        # 1. 替换 owner_type 的类型参数（如 tuple[T, Args...] 中的 {T}, {Args...}）
        # 这会替换模板中所有来自 owner_type 的类型参数占位符
        if type_info:
            emitter = TemplateEmitter()
            template = emitter._replace_type_parameters(template, type_info, context)

        # 2. 替换返回类型占位符（在模板参数替换之后，用实际推断的类型替换）
        if return_type_cpp_type:
            template = template.replace(
                return_type_temp_placeholder, return_type_cpp_type
            )

        return template

    def _apply_type_conversion(
        self,
        value_code: str,
        value_cpp_type: str | None,
        target_cpp_type: str,
        context: CodeGenContext,
    ) -> str:
        """应用类型转换（如果需要）

        Args:
            value_code: 值表达式代码
            value_cpp_type: 值的 C++ 类型
            target_cpp_type: 目标 C++ 类型
            context: 代码生成上下文

        Returns:
            str: 可能添加类型转换的代码
        """
        from mcpy.codegen.utils.type_conversion import apply_variant_conversion

        result = apply_variant_conversion(value_code, value_cpp_type, target_cpp_type)
        if result != value_code:
            context.require_include("mcpy/builtins.h")
        return result

    def _is_slice(self, slice_node: IRNode) -> bool:
        """检查节点是否是切片表达式

        Args:
            slice_node: 切片节点

        Returns:
            bool: 是否是切片
        """
        return isinstance(slice_node, IRSlice)

    def _emit_slice(
        self, node: IRSubscript, obj_code: str, context: CodeGenContext
    ) -> str | None:
        """生成切片访问代码

        Args:
            node: 下标访问节点
            obj_code: 对象表达式代码
            context: 代码生成上下文

        Returns:
            str | None: 生成的切片代码
        """
        generator = get_global_generator()

        if not isinstance(node.slice, IRSlice):
            return None

        slice_info = node.slice

        # 提取 start, stop, step
        # Python 中 None 表示使用默认值，在 C++ 中用 INT64_MIN 表示
        start_code = "INT64_MIN"
        stop_code = "INT64_MIN"
        step_code = "1"

        if slice_info.lower is not None:
            start_code = generator.generate_expression(slice_info.lower, context)

        if slice_info.upper is not None:
            stop_code = generator.generate_expression(slice_info.upper, context)

        if slice_info.step is not None:
            step_code = generator.generate_expression(slice_info.step, context)

        # 添加必要的包含文件
        context.require_include("mcpy/builtins.h")
        context.require_include("<cstdint>")

        # 生成切片调用代码
        code = f"mcpy::list_slice({obj_code}, {start_code}, {stop_code}, {step_code})"

        return code

    def get_priority(self) -> int:
        """下标访问发射器的优先级"""
        return 150  # 较高优先级，与属性访问和函数调用相同
